// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/connectors/connector.h>
#include <pollux/core/plan_node.h>
#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/testing/exec/util/hive_connector_test_base.h>
#include <pollux/parse/plan_nodeId_generator.h>

namespace fecebook::pollux::exec::test {
class IndexLookupJoinTestBase
    : public kumo::pollux::exec::test::HiveConnectorTestBase {
 protected:
  IndexLookupJoinTestBase() = default;

  struct SequenceTableData {
    kumo::pollux::RowVectorPtr keyData;
    kumo::pollux::RowVectorPtr valueData;
    kumo::pollux::RowVectorPtr tableData;
    std::vector<int64_t> minKeys;
    std::vector<int64_t> maxKeys;
  };

  static kumo::pollux::RowTypePtr concat(
      const kumo::pollux::RowTypePtr& a,
      const kumo::pollux::RowTypePtr& b);

  bool isFilter(const std::string& conditionSql) const;

  int getNumRows(const std::vector<int>& cardinalities);

  // Generate probe input for lookup join.
  // @param numBatches: number of probe batches.
  // @param batchSize: number of rows in each probe batch.
  // @param numDuplicateProbeRows: number of duplicates for each probe row so
  // the actual batch size is batchSize * numDuplicatesProbeRows.
  // @param tableData: contains the sequence table data including key vectors
  // and min/max key values.
  // @param probeJoinKeys: the prefix key colums used for equality joins.
  // @param inColumns: the ordered list of in conditions.
  // @param betweenColumns: the ordered list of between conditions.
  // @param equalMatchPct: percentage of rows in the probe input that matches
  // with the rows in index table.
  // @param betweenMatchPct: percentage of rows in the probe input that matches
  // the rows in index table with between conditions.
  // @param inMatchPct: percentage of rows in the probe input that matches the
  // rows in index table with in conditions.
  std::vector<kumo::pollux::RowVectorPtr> generateProbeInput(
      size_t numBatches,
      size_t batchSize,
      size_t numDuplicateProbeRows,
      SequenceTableData& tableData,
      std::shared_ptr<kumo::pollux::memory::MemoryPool>& pool,
      const std::vector<std::string>& probeJoinKeys,
      const std::vector<std::string> inColumns = {},
      const std::vector<std::pair<std::string, std::string>>& betweenColumns =
          {},
      std::optional<int> equalMatchPct = std::nullopt,
      std::optional<int> inMatchPct = std::nullopt,
      std::optional<int> betweenMatchPct = std::nullopt);

  // Makes lookup join plan with the following parameters:
  // @param indexScanNode: the index table scan node.
  // @param probeVectors: the probe input vectors.
  // @param leftKeys: the left join keys of index lookup join.
  // @param rightKeys: the right join keys of index lookup join.
  // @param joinType: the join type of index lookup join.
  // @param outputColumns: the output column names of index lookup join.
  // @param joinNodeId: returns the plan node id of the index lookup join
  // node.
  kumo::pollux::core::PlanNodePtr makeLookupPlan(
      const std::shared_ptr<kumo::pollux::core::PlanNodeIdGenerator>&
          planNodeIdGenerator,
      kumo::pollux::core::TableScanNodePtr indexScanNode,
      const std::vector<kumo::pollux::RowVectorPtr>& probeVectors,
      const std::vector<std::string>& leftKeys,
      const std::vector<std::string>& rightKeys,
      const std::vector<std::string>& joinConditions,
      kumo::pollux::core::JoinType joinType,
      const std::vector<std::string>& outputColumns,
      kumo::pollux::core::PlanNodeId& joinNodeId);

  void createDuckDbTable(
      const std::string& tableName,
      const std::vector<kumo::pollux::RowVectorPtr>& data);

  // Makes index table scan node with the specified index table handle.
  // @param outputType: the output schema of the index table scan node.
  // @param scanNodeId: returns the plan node id of the index table scan
  // node.
  kumo::pollux::core::TableScanNodePtr makeIndexScanNode(
      const std::shared_ptr<kumo::pollux::core::PlanNodeIdGenerator>&
          planNodeIdGenerator,
      const std::shared_ptr<kumo::pollux::connector::ConnectorTableHandle>
          indexTableHandle,
      const kumo::pollux::RowTypePtr& outputType,
      kumo::pollux::core::PlanNodeId& scanNodeId,
      std::unordered_map<
          std::string,
          std::shared_ptr<kumo::pollux::connector::ColumnHandle>>&
          assignments);

  // Generate sequence storage table which will be persisted by mock zippydb
  // client for testing.
  // @param keyCardinalities: specifies the number of unique keys per each index
  // column, which also determines the total number of rows stored in the
  // sequence storage table.
  // @param tableData: returns the sequence table data stats including the key
  // vector, value vector, table vector, and the min and max key values for each
  // index column.
  void generateIndexTableData(
      const std::vector<int>& keyCardinalities,
      SequenceTableData& tableData,
      std::shared_ptr<kumo::pollux::memory::MemoryPool>& pool);

  // Makes output schema from the index table scan node with the specified
  // column names.
  kumo::pollux::RowTypePtr makeScanOutputType(
      std::vector<std::string> outputNames);

  std::shared_ptr<kumo::pollux::exec::Task> runLookupQuery(
      const kumo::pollux::core::PlanNodePtr& plan,
      int numPrefetchBatches,
      const std::string& duckDbVefifySql);

  kumo::pollux::RowTypePtr keyType_;
  kumo::pollux::RowTypePtr valueType_;
  kumo::pollux::RowTypePtr tableType_;
  kumo::pollux::RowTypePtr probeType_;
  std::optional<kumo::pollux::RowTypePtr> partitionType_;
};
} // namespace fecebook::pollux::exec::test
